#include <cstdio>
#include <cublas_v2.h>
#include <cuda_runtime.h>

int main() {
    // 定义矩阵维度：A (2x3), B (3x4), C (2x4)
    const int m = 2, k = 3, n = 4;

    // 主机内存分配（行主序初始化）
    float h_A[m][k] = {{1, 2, 3}, {4, 5, 6}};
    float h_B[k][n] = {{1, 2, 3, 4}, {5, 6, 7, 8}, {9, 10, 11, 12}};
    float h_C[m][n] = {0};

    // 设备内存分配
    float *d_A, *d_B, *d_C;
    cudaMalloc(&d_A, m * k * sizeof(float));
    cudaMalloc(&d_B, k * n * sizeof(float));
    cudaMalloc(&d_C, m * n * sizeof(float));

    // 数据拷贝到设备（行主序 → cuBLAS 按列主序处理）
    cudaMemcpy(d_A, h_A, m * k * sizeof(float), cudaMemcpyHostToDevice);
    cudaMemcpy(d_B, h_B, k * n * sizeof(float), cudaMemcpyHostToDevice);

    // 创建 cuBLAS 句柄
    cublasHandle_t handle;
    cublasCreate(&handle);

    // 执行矩阵乘法 C = A * B
    const float alpha = 1.0f, beta = 0.0f;
    cublasSgemm(
        handle,
        CUBLAS_OP_N, CUBLAS_OP_N,   // 不转置 A 和 B
        n, m, k,                    // C 的维度为 m x n
        &alpha,
        d_B, n,                     // 左矩阵为B， n * k
        d_A, k,                     // 右矩阵为A， k * m
        &beta,
        d_C, n                      // C 是 n * m， 行数为m
    );

    // 拷贝结果回主机
    cudaMemcpy(h_C, d_C, m * n * sizeof(float), cudaMemcpyDeviceToHost);

    // 打印结果
    printf("C = \n");
    for (int i = 0; i < m; i++) {
        for (int j = 0; j < n; j++) {
            printf("%6.1f", h_C[i][j]);
        }
        printf("\n");
    }

    // 释放资源
    cublasDestroy(handle);
    cudaFree(d_A);
    cudaFree(d_B);
    cudaFree(d_C);
    return 0;
}